# 本文件主要包含 STO3G 轨道的构建及分子的构建

import numpy as np


class Gaussion_s:
    '''
    1s 高斯轨道
    '''
    center: float = 0.  # 中心坐标
    alpha: float = 1.   # 指数系数
    # 由于作业要求计算的都是线性分子，中心坐标只需要一个指标

    def __init__(self, center: float, alpha: float):
        self.center = center
        self.alpha = alpha


class STO3G:
    '''
    STO3G 轨道
        sto3g = d[0]*g[0] + d[1]*g[1] + d[2]*g[2]
    '''
    d: np.ndarray = np.zeros((3,))  # sto3g 展开系数
    g: list = [Gaussion_s(0, 1)]*3  # 1s 高斯轨道

    def __init__(self, d: np.ndarray, g: list):
        self.d = d
        self.g = g


class Atom:
    '''
    原子
    '''
    Z: int          # 原子序数，或者原子核电荷量
    R: float        # 位置坐标
    orbital: STO3G  # 轨道，由于只考虑 H 和 He，每个原子仅有一个轨道

    def __init__(self, Z: int, R: float, orbital: STO3G):
        self.Z = Z
        self.R = R
        self.orbital = orbital


class Molecule:
    '''
    分子
    '''
    name: str = ""  # 名字
    charge: int    # 电荷量
    atoms: list    # 含有的原子

    def __init__(self, name, charge, atoms):
        self.name = name
        self.charge = charge
        self.atoms = atoms


def hydrogen(center: float) -> Atom:
    '''
    构建氢原子
        center: 中心坐标
    '''
    return Atom(
        1, center,
        STO3G(
            np.array([0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00]),
            [
                Gaussion_s(center, 0.3425250914E+01),
                Gaussion_s(center, 0.6239137298E+00),
                Gaussion_s(center, 0.1688554040E+00)
            ]
        ))


def helium(center: float) -> Atom:
    '''
    构建氦原子
        center: 中心坐标
    '''
    return Atom(
        2, center,
        STO3G(
            np.array([0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00]),
            [
                Gaussion_s(center, 0.6362421394E+01),
                Gaussion_s(center, 0.1158922999E+01),
                Gaussion_s(center, 0.3136497915E+00)
            ]
        ))


def H2(r: float) -> Molecule:
    '''
    构建 H2 分子
    参数:
        r: 两原子之间的坐标
    '''
    return Molecule("H2", 0, [
        hydrogen(0.),
        hydrogen(r)
    ])


def HeH(r: float) -> Molecule:
    '''
    构建 HeH+ 离子
    参数:
        r: 两原子之间的坐标
    '''
    return Molecule("HeH+", 1, [
        helium(0.),
        hydrogen(r)
    ])
